#---------------------# 
# Read stratification #
#---------------------# 

#--------------------
# Import Modules
#--------------------

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from functions import PT  
from mpl_toolkits.axes_grid1 import ImageGrid
from scipy.integrate import cumtrapz 

#--------------------
# Constants 
#--------------------

G = 6.674e-11    # [G] = [N(m/kg)^2]
R = 13732.   
rs00 = 1.36206e9  # 1.1Myr
rs01 = 7.87093e8  # 6Myr
rs02 = 6.90278e8  # 9.25Myr
rs03 = 6.07751e8  # 14Myr

#-------------------------------------------------------
# Read Environmental Data with Polytropic aproximation  
#-------------------------------------------------------

#inputfile1 = '/home/guerrero/mfdyn/trunk/sms/data.txt'

inputfile00 = '/home/bonnie/ttauri/ttauri00/data_env.txt'   # (m = 1.4995)
inputfile01 = '/home/bonnie/ttauri/ttauri01/data_env.txt'   # (m = 1.4995)
inputfile02 = '/home/bonnie/ttauri/ttauri02/data_env.txt'   # (m = 1.4998, 2.0)
inputfile03 = '/home/bonnie/ttauri/ttauri03/data_env.txt'   # (m = 1.4998, 1.65)

f_env00 = open(inputfile00, 'r')
files_env00 = np.genfromtxt(f_env00, usecols = (0, 1, 2, 3))
r_env00 = files_env00[:,0]     # [r] = [m]
T_env00 = files_env00[:,1]     # [T_e] = [K]                  
rho_env00 = files_env00[:,2]   # [rho_e] = [Kg/m^3]  
pt_env00 = files_env00[:,3]
f_env00.close()

f_env01 = open(inputfile01, 'r')
files_env01 = np.genfromtxt(f_env01, usecols = (0, 1, 2, 3))
r_env01 = files_env01[:,0]     # [r] = [m]
T_env01 = files_env01[:,1]     # [T_e] = [K]                  
rho_env01 = files_env01[:,2]   # [rho_e] = [Kg/m^3]  
pt_env01 = files_env01[:,3]
f_env01.close()

f_env02 = open(inputfile02, 'r')
files_env02 = np.genfromtxt(f_env02, usecols = (0, 1, 2, 3))
r_env02 = files_env02[:,0]     # [r] = [m]
T_env02 = files_env02[:,1]     # [T_e] = [K]                  
rho_env02 = files_env02[:,2]   # [rho_e] = [Kg/m^3]  
pt_env02 = files_env02[:,3]
f_env02.close()

f_env03 = open(inputfile03, 'r')
files_env03 = np.genfromtxt(f_env03, usecols = (0, 1, 2, 3))
r_env03 = files_env03[:,0]     # [r] = [m]
T_env03 = files_env03[:,1]     # [T_e] = [K]                  
rho_env03 = files_env03[:,2]   # [rho_e] = [Kg/m^3]  
pt_env03 = files_env03[:,3]
f_env03.close()

rr = np.linspace(0.1, 0.95, 256)
r00 = rr*rs00
r01 = rr*rs01
r02 = rr*rs02
r03 = rr*rs03

#-------------------------------------------------------
# Read Isentropic Data with Polytropic aproximation
#-------------------------------------------------------

#input00 ='/home/guerrero/mfdyn/trunk/sms/data.txt'   

input00 ='/home/bonnie/ttauri/ttauri00/data_isen.txt'
input01 ='/home/bonnie/ttauri/ttauri01/data_isen.txt'
input02 ='/home/bonnie/ttauri/ttauri02/data_isen.txt'
input03 ='/home/bonnie/ttauri/ttauri03/data_isen.txt'

f_isen00 = open(input00, 'r')
files_isen00 = np.genfromtxt(f_isen00, usecols = (0, 1, 2, 3))
r_isen00 = files_isen00[:,0]     # [r] = [m]
T_isen00 = files_isen00[:,1]     # [T] = [K]                  
rho_isen00 = files_isen00[:,2]   # [rho] = [Kg/m^3]  
pt_isen00 = files_isen00[:,3]
f_isen00.close()

f_isen01 = open(input01, 'r')
files_isen01 = np.genfromtxt(f_isen01, usecols = (0, 1, 2, 3))
r_isen01 = files_isen01[:,0]     # [r] = [m]
T_isen01 = files_isen01[:,1]     # [T] = [K]                  
rho_isen01 = files_isen01[:,2]   # [rho] = [Kg/m^3]  
pt_isen01 = files_isen01[:,3]
f_isen01.close()

f_isen02 = open(input02, 'r')
files_isen02 = np.genfromtxt(f_isen02, usecols = (0, 1, 2, 3))
r_isen02 = files_isen02[:,0]     # [r] = [m]
T_isen02 = files_isen02[:,1]     # [T] = [K]                  
rho_isen02 = files_isen02[:,2]   # [rho] = [Kg/m^3]  
pt_isen02 = files_isen02[:,3]
f_isen02.close()

f_isen03 = open(input03, 'r')
files_isen03 = np.genfromtxt(f_isen03, usecols = (0, 1, 2, 3))
r_isen03 = files_isen03[:,0]     # [r] = [m]
T_isen03 = files_isen03[:,1]     # [T] = [K]                  
rho_isen03 = files_isen03[:,2]   # [rho] = [Kg/m^3]  
pt_isen03 = files_isen03[:,3]
f_isen03.close()


#------------------------------------------
# Read Data Stelar Model
#------------------------------------------

inputfile00 ='/home/bonnie/ttauri/ttauri00/stratification_0.61Msun_1.1Myr.txt'
inputfile01 ='/home/bonnie/ttauri/ttauri01/stratification_0.61Msun_6Myr.txt'
inputfile02 ='/home/bonnie/ttauri/ttauri02/stratification_0.61Msun_9.25Myr.txt'
inputfile03 ='/home/bonnie/ttauri/ttauri03/stratification_0.61Msun_14Myr.txt'

f00 = open(inputfile00, 'r')
files00 = np.genfromtxt(f00, usecols = (0,1,2,3))
r_sm00 = 1e-2*files00[:,0]  # [r] = [m]
P00 = files00[:,1]  
T00 = files00[:,2]          # [T] = [K]
rho00 = 1e3*files00[:,3]    # [rho] = [Kg/m^3]
f00.close()

f01 = open(inputfile01, 'r')
files01 = np.genfromtxt(f01, usecols = (0,1,2,3))
r_sm01 = 1e-2*files01[:,0]  # [r] = [m]
P01 = files01[:,1]  
T01 = files01[:,2]          # [T] = [K]
rho01 = 1e3*files01[:,3]    # [rho] = [Kg/m^3]
f01.close()

f02 = open(inputfile02, 'r')
files02 = np.genfromtxt(f02, usecols = (0,1,2,3))
r_sm02 = 1e-2*files02[:,0]  # [r] = [m]
P02 = files02[:,1]  
T02 = files02[:,2]          # [T] = [K]
rho02 = 1e3*files02[:,3]    # [rho] = [Kg/m^3]
f02.close()

f03 = open(inputfile03, 'r')
files03 = np.genfromtxt(f03, usecols = (0,1,2,3))
r_sm03 = 1e-2*files03[:,0]  # [r] = [m]
P03 = files03[:,1]  
T03 = files03[:,2]          # [T] = [K]
rho03 = 1e3*files03[:,3]    # [rho] = [Kg/m^3]
f03.close()

M00 = 4 * np.pi * cumtrapz(rho00*r_sm00**2, r_sm00, initial = 0)
M01 = 4 * np.pi * cumtrapz(rho01*r_sm01**2, r_sm01, initial = 0)
M02 = 4 * np.pi * cumtrapz(rho02*r_sm02**2, r_sm02, initial = 0)
M03 = 4 * np.pi * cumtrapz(rho03*r_sm03**2, r_sm03, initial = 0)

## These arrays have points for the same radius used in polytropic model:

P00 = np.interp(r00, r_sm00, P00)
T00 = np.interp(r00, r_sm00, T00)
rho00 = np.interp(r00, r_sm00, rho00)
M00 = np.interp(r00, r_sm00, M00)
g00 = G * M00 / r00**2
Tb00 = T00[0]
rhob00 = rho00[0]
pt00 = T00*(Tb00*rhob00/(T00*rho00))**(2./5.)

P01 = np.interp(r01, r_sm01, P01)
T01 = np.interp(r01, r_sm01, T01)
rho01 = np.interp(r01, r_sm01, rho01)
M01 = np.interp(r01, r_sm01, M01)
g01 = G * M01 / r01**2
Tb01 = T01[0]
rhob01 = rho01[0]
pt01 = T01*(Tb01*rhob01/(T01*rho01))**(2./5.)

P02 = np.interp(r02, r_sm02, P02)
T02 = np.interp(r02, r_sm02, T02)
rho02 = np.interp(r02, r_sm02, rho02)
M02 = np.interp(r02, r_sm02, M02)
g02 = G * M02 / r02**2
Tb02 = T02[0]
rhob02 = rho02[0]
pt02 = T02*(Tb02*rhob02/(T02*rho02))**(2./5.)

P03 = np.interp(r03, r_sm03, P03)
T03 = np.interp(r03, r_sm03, T03)
rho03 = np.interp(r03, r_sm03, rho03)
M03 = np.interp(r03, r_sm03, M03)
g03 = G * M03 / r03**2
Tb03 = T03[0]
rhob03 = rho03[0]
pt03 = T03*(Tb03*rhob03/(T03*rho03))**(2./5.)

#--------------------------------------------
# Polynomial approximation for Gravity 
#--------------------------------------------

## FOR STELLAR MODEL:

coefficients00 = np.polyfit(r00, g00, 3)
polynomial00 = np.poly1d(coefficients00)
g_p00 = polynomial00(r00)
print g_p00[0]
g_p00 = 26.*g_p00/g_p00[0]

coefficients01 = np.polyfit(r01, g01, 3)
polynomial01 = np.poly1d(coefficients01)
g_p01 = polynomial01(r01)
print g_p01[0]

coefficients02 = np.polyfit(r02, g02, 3)
polynomial02 = np.poly1d(coefficients02)
g_p02 = polynomial02(r02)
print g_p02[0]

coefficients03 = np.polyfit(r03, g03, 3)
polynomial03 = np.poly1d(coefficients03)
g_p03 = polynomial03(r03)
print g_p03[0]
g_p03 = 132.5*g_p03/g_p03[0]

#print polynomial
#print 'g = 26.*(', polynomial00/g_p00[0], ')'
#print 'rho =', rhob00
#print 'T =', Tb00
#print 'gb =', g_p00[0]

#-----------------------------------------
# Plot
#-----------------------------------------

font = {'family' : 'serif',
        'color'  : 'black',
        'fontweight' : 'normal',
        'fontsize'   : 15,
        'horizontalalignment' : 'center',
        'verticalalignment': 'center',
       }

font1 = {'family' : 'serif'}

plt.rc('text', usetex=True)

plt.plot(rr, g00, 'm', rr, g_p00, 'm--', linewidth = 1.5)
plt.title(r'$TT00$', fontdict = font)
plt.legend((r'$g_{\rm sm}$', r'$g$'), loc = 8, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('linear')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'both', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$g \left [\rm m/s^{2} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

plt.plot(rr, g03, 'b', rr, g_p03, 'b--', linewidth = 1.5)
plt.title(r'$TT03$', fontdict = font)
plt.legend((r'$g_{\rm sm}$', r'$g$'), loc = 8, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('linear')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'both', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$g \left [\rm m/s^{2} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

plt.plot(rr, g_p00, 'm', rr, g00, 'm:', rr, g_p03, 'b', rr, g03, 'b:', linewidth = 1.5)
plt.legend((r'$g_{\rm FIT}$', r'$g_{\rm sm}$', r'$g_{\rm FIT}$', r'$g_{\rm sm}$'), loc = 8, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('linear')
plt.xlim(rr.min(),rr.max())
plt.ylim(0,1.01*g_p03.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'both', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$g \left [\rm m/s^{2} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

plt.plot( rr, T_env00, 'm--', rr, T_isen00, 'm:', rr, T00, 'm', linewidth = 1.5)
plt.title('TT00', fontdict = font)
plt.legend(( r'$T_{\rm e}$', r'$T_{\rm s}$', r'$T_{\rm sm}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('log')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'x', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$T[\rm K]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

plt.plot( rr, T_env03, 'b--', rr, T_isen03, 'b:', rr, T03, 'b', linewidth = 1.5)
plt.title('TT03', fontdict = font)
plt.legend(( r'$T_{\rm e}$', r'$T_{\rm s}$', r'$T_{\rm sm}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('log')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'x', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$T[\rm K]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

plt.plot(rr, T_env00, 'm--', rr, T_isen00, 'm:', rr, T00, 'm', rr, T_env03, 'b--', rr, T_isen03, 'b:', rr, T03, 'b',linewidth = 1.5)
plt.legend((r'$T_{\rm e}$', r'$T_{\rm s}$', r'$T_{\rm sm}$', r'$T_{\rm e}$', r'$T_{\rm s}$', r'$T_{\rm sm}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('log')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'x', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$T[\rm K]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()
 
plt.plot( rr, rho_env00, 'm--', rr, rho_isen00, 'm:', rr, rho00, 'm', linewidth = 1.5)
plt.title('TT00', fontdict = font)
plt.legend(( r'$\rho_{\rm e}$', r'$\rho_{\rm s}$', r'$\rho_{\rm sm}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('log')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'x', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$\rho \left [\rm Kg/m^{3} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

plt.plot(rr, rho_env03, 'b--', rr, rho_isen03, 'b:', rr, rho03, 'b', linewidth = 1.5)
plt.title('TT03', fontdict = font)
plt.legend((r'$\rho_{\rm e}$', r'$\rho_{\rm s}$', r'$\rho_{\rm sm}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('log')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'x', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$\rho \left [\rm Kg/m^{3} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()

         
plt.plot(rr, rho_env00, 'm--', rr, rho_isen00, 'm:', rr, rho00, 'm', rr, rho_env03, 'b--', rr, rho_isen03, 'b:', rr, rho03, 'b', linewidth = 1.5)
plt.legend((r'$\rho_{\rm e}$', r'$\rho_{\rm s}$', r'$\rho_{\rm sm}$', r'$\rho_{\rm e}$', r'$\rho_{\rm s}$', r'$\rho_{\rm sm}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('log')
plt.xlim(rr.min(),rr.max())
plt.ticklabel_format(style = 'scientific', scilimits = (-2,2), axis = 'x', pad = 10)
plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$\rho \left [\rm Kg/m^{3} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()
